
import json
import argparse
import pandas as pd
from multiprocessing import Pool



# ---------------------大模型体系相关模块---------------------------------------
from zhipuai import ZhipuAI
# ChatOpenAI 是专门设计用于与 ChatGPT 或其他 OpenAI 模型进行交互的
from langchain_openai import ChatOpenAI   
from langchain.prompts import (
    # 这是一个通用的聊天提示模板,用于定义整个聊天交互的结构
    # 它可以包含系统消息、用户消息、AI消息以及任何其他相关的元素
    ChatPromptTemplate,  
    # MessagesPlaceholder可以作为一个占位符,在ChatPromptTemplate或其他模板中使用
    # 以便在运行时插入动态的消息列表,这增加了模板的灵活性和可定制性
    MessagesPlaceholder, 
    # 专门用于定义系统级别的消息,这些消息通常不是由用户或AI直接发出的
    # 而是由系统生成,用于提供额外的上下文或指导
    SystemMessagePromptTemplate,
    # 人工消息提示模板,这是用户发送的消息
    HumanMessagePromptTemplate,  
    # AI 消息提示模板,这是从 AI 发送的消息
    # 用于定义AI（通常是聊天模型）在对话中发送的消息格式
    AIMessagePromptTemplate,     
)
# OutputParser 类解析 LLM 调用的输出:
from langchain_core.output_parsers import JsonOutputParser



############################### Prompt Template ################################
def prompt_generate(system, ai):
    """ 
        生成链中需要的Prompt模板
    """
    prompt = ChatPromptTemplate(
        messages=[
            # 定义一个系统消息
            SystemMessagePromptTemplate.from_template(system), 
            # 定义一个AI消息     
            AIMessagePromptTemplate.from_template(ai),             
            # 定义一个人类消息     
            HumanMessagePromptTemplate.from_template("{topic}")  
        ]
    )
    return prompt



################################### Chains #####################################
"""
    模型温度:
        低温度(例如0.1):模型可能会生成非常标准、常见的回答,缺乏新意
        中温度(例如0.5-0.9):
            模型通常能在准确性和创造性之间找到一个良好的平衡,适用于大多数情况
        高温度(例如0.95或更高):
            模型可能会生成更加独特、有趣的回答,但也可能包含更多的错误或不连贯之处
"""



def chat_glm(prompt, user, parser):
    """ 
        输入提示词,运行glm-4-flash模型的一个链,然后格式化输出
    """
    # 创建一个ChatOpenAI实例
    llm = ChatOpenAI(
        # 设置模型的"温度"参数，影响生成内容的随机性和创造性
        temperature=0.5,   
        # 设置使用的模型
        model="glm-4-flash",      
        # 这里API Key
        openai_api_key="e7794ecfd23516cd297dc83e7f57a0e9.1ta3BPaDit6LuXJj",   
        # 这里API Base
        openai_api_base="https://open.bigmodel.cn/api/paas/v4/" 
    )
    # 定义个langchain的链
    chain = prompt | llm | parser
    # 调用链，生成链的结果
    response = chain.invoke({"topic": user})

    return response



##################################### Main #####################################
def gene_cluster_chatglm4(prompt, user_prompt, jsonparser, cluster_name):
    break_number = 1
    input_summary = ""
    while True:  # 如果结果报错，那么就重新用大模型生成，直到不报错为止
        if break_number == 5:
            break
        try:
            input_summary = chat_glm(
                prompt, user_prompt, jsonparser
            )
            # print(f"\n\n{cluster_name}:\n\t",input_summary)
            print(f"\n\n{cluster_name} 已经运行完成！\n")
            break
        except:
            pass
        # if task_num == 50:
        #     break 
    row = [cluster_name, input_summary["main_category"],  input_summary["summary_feature"]]
    
    return row 



def multi_summary(promt_file, output_path, multi):
    with open(promt_file) as file:
        prompt_dict = json.load(file)
    # 创建结果解析器
    jsonparser = JsonOutputParser()
    system = [
        "As a bioinformatics expert, you will, based on the gene cluster information provided by the user,",
        "primarily identify the metabolic pathway or metabolite associated with this gene set, without including the word 'Cluster' in its name.",
        "The functional feature labels of gene clusters need to be indivisible concepts,",
        "And summarize the gene cluster by giving it a general name, which represents its classification"
    ]
    ai = [
        "The output is in JSON format,The value of 'summary_feature:' should be enclosed in brackets even if it contains only one element.", 
        "main_category: Please fill in a general name for the gene cluster here, with a length of several to a dozen words. ",
        "The name should not contain the word 'cluster'. You can fill in the main metabolic pathway of the cluster or the main metabolite in the cluster",
        "summary_feature: [",
            "Please fill in the functional feature label 1 for the gene cluster with a brief description. The response should be in English.",
            "Please fill in the functional feature label 2 for the gene cluster with a brief description. The response should be in English.",
            "Please fill in the functional feature label 3 for the gene cluster with a brief description. The response should be in English.",
        "]"
    ]
    tasks = []
    for cluster_prompt in prompt_dict:
        cluster = list(cluster_prompt.keys())[0]
        cluster_name = f"cluster_{cluster}"
        gene_prompt = "In the gene cluster, the description categories and the number of genes are as follows in a two-dimensional list:" + str(cluster_prompt[cluster]["gene_prompt"])
        pfam_prompt = "In the gene cluster, the PFAM descriptions of the genes are as follows in a two-dimensional list:" + str(cluster_prompt[cluster]["pfam_prompt"])
        gog_category = "In the gene cluster, the COG classifications of the genes are as follows in a two-dimensional list:" + str(cluster_prompt[cluster]["gog_prompt"])
        user_prompt = (
            "The gene cluster name input by the user is:" + cluster_name + "\n\n" +
            gene_prompt + "\n\n" +
            pfam_prompt + "\n\n" +
            gog_category + "\n\n" +
            "Based on the above information, please summarize the metabolic pathway or the primary metabolic compound that this gene cluster is primarily associated with, and provide a functional feature label." +
            "Do not use similar phrases'Signal Transduction and Ion Transport Metabolic Cluster' " + 
            "'ATP Synthesis and Transcription Regulation'," +
            "Using 'and' to link two concepts should be avoided; the two concepts should be separated." +
            "It is given in the 'main_category':" +
            "'NAC-Arv1 Protein Complex' 'Auxin Responsive Gene Module' 'Zinc Finger Regulation'" +
            "Such a concise and accurate answer." +  
            "At the same time, summarize the main aspects of the functional characteristics of this gene cluster and fill them in 'summary_feature'," +
            "'Unknown functional' or 'unknown' in the input and output needs to be filtered out." +
            "And output it in the generated JSON format."
        )
        prompt = prompt_generate(
            " ".join(system),
            " ".join(ai))
        tasks.append((prompt, user_prompt, jsonparser, cluster_name))
    with Pool(int(multi)) as pool:
        result = pool.starmap(gene_cluster_chatglm4, tasks)
    df = pd.DataFrame(result, columns=["cluster", "main_category", "summary_feature"])
    df.to_excel(output_path, index=False)



def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("-input","--input_path")
    parser.add_argument("-output","--output_path")
    parser.add_argument("-multi","--multi")
    args = parser.parse_args()
    
    return args



if __name__ == "__main__":
    params = parse_args()
    promt_file = params.input_path
    output_path = params.output_path
    multi = params.multi
    # promt_file = "Aco2.json"
    # output_path = "Aco2_summary.xlsx"
    # multi = 50
    multi_summary(promt_file, output_path, multi)